﻿#include "json.hpp"
#include <filesystem>
#include <fstream>
#include <string>
#include <unordered_map>

void LabelMeConvertDarknet(const std::string &labelme_path, const std::string &darknet_path)
{
    namespace fs          = std::filesystem;
    using json            = nlohmann::json;
    auto labels_directory = darknet_path + "/labels";
    auto train_path       = darknet_path + "/train.txt";
    auto dataset_path     = darknet_path + "/dataset.txt";
    fs::create_directories(labels_directory);

    std::ofstream ofs_train(train_path);
    std::ofstream ofs_dataset(dataset_path);

    std::unordered_map<std::string, int> label_ids;
    std::vector<std::string> names;
    for (const auto &entry : fs::directory_iterator(labelme_path))
    {
        if (entry.is_directory())
            continue;

        const auto file_path = entry.path();
        const auto file_ext  = file_path.extension().string();

        if (file_ext != ".json")
            continue;

        std::string label_path = labels_directory + "/" + file_path.filename().replace_extension(".txt").string();
        std::ifstream ifs(file_path);
        std::ofstream ofs(label_path);

        json data              = json::parse(ifs);
        int image_width        = data["imageWidth"];
        int image_height       = data["imageHeight"];
        std::string image_path = data["imagePath"];
        image_path = labelme_path + "/" + image_path;
        ofs_dataset << image_path << std::endl;
        for (auto &shape : data["shapes"])
        {
            std::string label = shape["label"];
            if (label_ids.find(label) == label_ids.end())
            {
                label_ids[label] = names.size();
                names.push_back(label);
            }
            int id       = label_ids[label];
            float left   = shape["points"][0][0];
            float top    = shape["points"][0][1];
            float right  = shape["points"][1][0];
            float bottom = shape["points"][1][1];
            float s_x    = (left + right) / 2.0f / image_width;
            float s_y    = (top + bottom) / 2.0f / image_height;
            float s_w    = (right - left) / image_width;
            float s_h    = (bottom - top) / image_height;
            // id center_x center_y width height
            ofs << id << " " << s_x << " " << s_y << " " << s_w << " " << s_h << std::endl;
            ofs_dataset << s_x << " " << s_y << " " << s_w << " " << s_h << " " << id << " ";
        }
        ofs_dataset << std::endl;
        ofs_train << image_path << std::endl;
    }

    auto names_path = darknet_path + "/names";
    std::ofstream ofs(names_path);
    for (auto name : names)
    {
        ofs << name << std::endl;
    }
}

int main(int argc, char *argv[])
{
    LabelMeConvertDarknet("D:/workspace/cpp/darknet/data/qrcode/images", "D:/workspace/cpp/darknet/data/qrcode");
}